import torch


class SwishGluFeedForward(torch.nn.Module):
    def __init__(self, *,
        i_size: int | None = None,
        h_size: int,
        o_size: int,
        bias: bool = True,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        self.i1_linear = torch.nn.Linear(i_size, h_size, bias=bias, dtype=dtype, device=device) \
            if i_size is not None else torch.nn.LazyLinear(h_size, bias=bias, dtype=dtype, device=device)
        self.i2_linear = torch.nn.Linear(i_size, h_size, bias=bias, dtype=dtype, device=device) \
            if i_size is not None else torch.nn.LazyLinear(h_size, bias=bias, dtype=dtype, device=device)
        self.o_linear = torch.nn.Linear(h_size, o_size, bias=bias, dtype=dtype, device=device)

    def forward(self,
        x: torch.Tensor, *,
        dropout: float | None = None,
    ) -> torch.Tensor:
        """
        :param x: shape=[..., i_size]
        :param dropout: float
        :return: shape=[..., o_size]
        """
        x1 = self.i1_linear(x)
        # [..., h_size]

        x2 = self.i2_linear(x)
        # [..., h_size]

        x = x1 * torch.nn.functional.silu(x2)
        # [..., h_size]

        if dropout is not None:
            x = torch.nn.functional.dropout(x, dropout)
            # [..., h_size]

        x = self.o_linear(x)
        # [..., o_size]

        return x
